
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>

#include <lua.h>
#include <lualib.h>
#include <lauxlib.h>

#define BUFFSIZE 1024

static int my_encrypt(lua_State *L);
static int my_pub_encrypt(lua_State *L);
static int my_decrypt(lua_State *L);

#include <openssl/evp.h>
#include <openssl/pem.h>

static int my_encrypt(lua_State *L) {
  EVP_PKEY *pkey = NULL;
  FILE *file = NULL;
  EVP_PKEY_CTX *ctx = NULL;
  unsigned char *p_en = NULL;
  size_t flen, outlen;
  const char *str = (const char *)luaL_checklstring(L, 1, &flen);
  const char *path_key = "rsa_pub.key";

  if ((file = fopen(path_key, "r")) == NULL) {
    perror("open key file error");
    goto err;
  }

  if ((pkey = PEM_read_PUBKEY(file, NULL, NULL, NULL)) == NULL) {
    ERR_print_errors_fp(stdout);
    goto err;
  }

  ctx = EVP_PKEY_CTX_new(pkey, NULL);
  if (!ctx || EVP_PKEY_encrypt_init(ctx) <= 0) {
    goto err;
  }

  if (EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_NO_PADDING) <= 0) {
    goto err;
  }

  if (EVP_PKEY_encrypt(ctx, NULL, &outlen, (const unsigned char *)str, flen) <=
      0) {
    goto err;
  }

  p_en = (unsigned char *)malloc(outlen + 1);
  if (!p_en) {
    goto err;
  }
  memset(p_en, 0, outlen + 1);

  if (EVP_PKEY_encrypt(ctx, p_en, &outlen, (const unsigned char *)str, flen) <=
      0) {
    goto err;
  }

  lua_pushlstring(L, (const char *)p_en, outlen);

  free(p_en);
  EVP_PKEY_CTX_free(ctx);
  EVP_PKEY_free(pkey);
  fclose(file);
  return 1;

err:
  if (p_en)
    free(p_en);
  if (ctx)
    EVP_PKEY_CTX_free(ctx);
  if (pkey)
    EVP_PKEY_free(pkey);
  if (file)
    fclose(file);
  return 0;
}

static int my_pub_encrypt(lua_State *L) {
  unsigned char *p_en = NULL;
  EVP_PKEY *pkey = NULL;
  BIO *bp = NULL;
  EVP_PKEY_CTX *ctx = NULL;
  size_t flen, plen, outlen;
  const char *str = (const char *)luaL_checklstring(L, 1, &flen);
  const char *pub_key = (const char *)luaL_checklstring(L, 2, &plen);

  if ((bp = BIO_new_mem_buf((void *)pub_key, -1)) == NULL) {
    printf("BIO_new_mem_buf failed!\n");
    goto err;
  }

  if ((pkey = PEM_read_bio_PUBKEY(bp, NULL, NULL, NULL)) == NULL) {
    ERR_print_errors_fp(stdout);
    goto err;
  }

  if ((ctx = EVP_PKEY_CTX_new(pkey, NULL)) == NULL ||
      EVP_PKEY_encrypt_init(ctx) <= 0) {
    goto err;
  }

  if (EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_NO_PADDING) <= 0) {
    goto err;
  }

  if (EVP_PKEY_encrypt(ctx, NULL, &outlen, (const unsigned char *)str, flen) <=
      0) {
    goto err;
  }

  p_en = (unsigned char *)malloc(outlen + 1);
  if (!p_en) {
    goto err;
  }
  memset(p_en, 0, outlen + 1);

  if (EVP_PKEY_encrypt(ctx, p_en, &outlen, (const unsigned char *)str, flen) <=
      0) {
    goto err;
  }

  lua_pushlstring(L, (const char *)p_en, outlen);

  free(p_en);
  EVP_PKEY_CTX_free(ctx);
  EVP_PKEY_free(pkey);
  BIO_free_all(bp);
  return 1;

err:
  if (p_en)
    free(p_en);
  if (ctx)
    EVP_PKEY_CTX_free(ctx);
  if (pkey)
    EVP_PKEY_free(pkey);
  if (bp)
    BIO_free_all(bp);
  return 0;
}

static int my_decrypt(lua_State *L) {
  EVP_PKEY *pkey = NULL;
  FILE *file = NULL;
  EVP_PKEY_CTX *ctx = NULL;
  unsigned char *p_de = NULL;
  size_t flen, outlen;
  const char *str = (const char *)luaL_checklstring(L, 1, &flen);
  const char *path_key = "rsa.key";

  if ((file = fopen(path_key, "r")) == NULL) {
    perror("open key file error");
    goto err;
  }

  if ((pkey = PEM_read_PrivateKey(file, NULL, NULL, NULL)) == NULL) {
    ERR_print_errors_fp(stdout);
    goto err;
  }

  if ((ctx = EVP_PKEY_CTX_new(pkey, NULL)) == NULL ||
      EVP_PKEY_decrypt_init(ctx) <= 0) {
    goto err;
  }

  if (EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_NO_PADDING) <= 0) {
    goto err;
  }

  if (EVP_PKEY_decrypt(ctx, NULL, &outlen, (const unsigned char *)str, flen) <=
      0) {
    goto err;
  }

  p_de = (unsigned char *)malloc(outlen + 1);
  if (!p_de) {
    goto err;
  }
  memset(p_de, 0, outlen + 1);

  if (EVP_PKEY_decrypt(ctx, p_de, &outlen, (const unsigned char *)str, flen) <=
      0) {
    goto err;
  }

  lua_pushlstring(L, (const char *)p_de, outlen);

  free(p_de);
  EVP_PKEY_CTX_free(ctx);
  EVP_PKEY_free(pkey);
  fclose(file);
  return 1;

err:
  if (p_de)
    free(p_de);
  if (ctx)
    EVP_PKEY_CTX_free(ctx);
  if (pkey)
    EVP_PKEY_free(pkey);
  if (file)
    fclose(file);
  return 0;
}

static const luaL_Reg rsa_funcs[] = {{"encrypt", my_encrypt},
                                     {"pub_encrypt", my_pub_encrypt},
                                     {"decrypt", my_decrypt},
                                     {NULL, NULL}};

/* 模块注册函数 */
LUALIB_API int luaopen_rsa(lua_State *L) {
  luaL_newlib(L, rsa_funcs);
  return 1;
}
